import os
import matplotlib.pyplot as plt
import numpy as np

# Create output directory
OUT_DIR = "Pictures"
os.makedirs(OUT_DIR, exist_ok=True)
OUT_PNG1 = os.path.join(OUT_DIR, "rule_total_evolution.png")
OUT_PNG2 = os.path.join(OUT_DIR, "rule_category_distribution.png")

# Experimental data
rounds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, "...", "Final"]
total_rules = [0, 22, 33, 44, 52, 57, 65, 67, 74, 80, 85, None, 87]
IG = [0, 3, 6, 7, 8, 11, 12, 12, 12, 13, 13, None, 14]
DT = [0, 8, 13, 19, 21, 21, 26, 26, 30, 31, 31, None, 20]
DE = [0, 0, 1, 3, 4, 4, 5, 5, 6, 8, 8, None, 7]
PE = [0, 10, 12, 14, 18, 20, 21, 23, 25, 27, 31, None, 42]
SP = [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, None, 4]

# Handle x-axis positions for omitted rounds
plot_rounds = list(range(11)) + [13]  # Only positions with data
plot_total = total_rules[:11] + [total_rules[-1]]
plot_IG = IG[:11] + [IG[-1]]
plot_DT = DT[:11] + [DT[-1]]
plot_DE = DE[:11] + [DE[-1]]
plot_PE = PE[:11] + [PE[-1]]
plot_SP = SP[:11] + [SP[-1]]

# Set x-axis labels
xtick_positions = list(range(11)) + [13]
xtick_labels = [f'Round {i}' if i > 0 else 'Initial' for i in range(11)] + ['Final']

# Figure 1: Total rule count evolution
fig1, ax1 = plt.subplots(1, 1, figsize=(10, 6))

ax1.plot(plot_rounds, plot_total, 'o-', linewidth=2, markersize=6, color='#2E86AB')
# Add omission line
ax1.plot([10, 13], [85, 87], '--', color='gray', alpha=0.7, linewidth=1)
ax1.axvspan(10.5, 12.5, alpha=0.1, color='gray')
ax1.text(11.5, 86, '...', fontsize=12, ha='center', va='center', 
         bbox=dict(boxstyle="round,pad=0.3", facecolor='white', alpha=0.8))

ax1.set_ylabel('Total Number of Rules', fontsize=12)
ax1.set_title('LLM Rule Generation: Total Rule Count Evolution', fontsize=14, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.set_ylim(0, 95)

ax1.set_xticks(xtick_positions)
ax1.set_xticklabels(xtick_labels, rotation=45, ha='right')

plt.tight_layout()
plt.savefig(OUT_PNG1, dpi=300, bbox_inches='tight')
print(f"Figure 1 saved to: {OUT_PNG1}")
plt.close()

# Figure 2: Rule category distribution evolution
fig2, ax2 = plt.subplots(1, 1, figsize=(12, 6))

width = 0.8
ax2.bar(plot_rounds, plot_IG, width, label='IG (Information Gathering)', color='#F18F01')
ax2.bar(plot_rounds, plot_DT, width, bottom=np.array(plot_IG), 
        label='DT (Data Transmission)', color='#C73E1D')
ax2.bar(plot_rounds, plot_DE, width, 
        bottom=np.array(plot_IG)+np.array(plot_DT), 
        label='DE (Data Encoding)', color='#593E2E')
ax2.bar(plot_rounds, plot_PE, width, 
        bottom=np.array(plot_IG)+np.array(plot_DT)+np.array(plot_DE), 
        label='PE (Payload Execution)', color='#176582')
ax2.bar(plot_rounds, plot_SP, width, 
        bottom=np.array(plot_IG)+np.array(plot_DT)+np.array(plot_DE)+np.array(plot_PE), 
        label='SP (Special Patterns)', color='#8E44AD')

# Add omission area indication
ax2.axvspan(10.5, 12.5, alpha=0.1, color='gray')
ax2.text(11.5, 45, '...', fontsize=12, ha='center', va='center',
         bbox=dict(boxstyle="round,pad=0.3", facecolor='white', alpha=0.8))

ax2.set_ylabel('Number of Rules by Category', fontsize=12)
ax2.set_title('Rule Category Distribution Evolution', fontsize=14, fontweight='bold')
ax2.legend(loc='upper left', bbox_to_anchor=(1.02, 1.0))
ax2.grid(True, alpha=0.3)
ax2.set_ylim(0, 95)

ax2.set_xticks(xtick_positions)
ax2.set_xticklabels(xtick_labels, rotation=45, ha='right')

plt.tight_layout()
plt.savefig(OUT_PNG2, dpi=300, bbox_inches='tight')
print(f"Figure 2 saved to: {OUT_PNG2}")
plt.close()

print("Both figures generated successfully!")